package com.xiashitech.agent.instrumentation.undertow.handler.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.config.StartConfig;
import com.xiashitech.agent.instrumentation.base.context.ReentryContextStorage;
import com.xiashitech.agent.instrumentation.base.context.XSBaggage;
import com.xiashitech.agent.instrumentation.base.propagator.W3CReentryBaggagePropagator;
import com.xiashitech.agent.instrumentation.bridge.Java8Bridge;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSUndertowHandlerCallDepth;
import com.xiashitech.agent.instrumentation.undertow.common.UndertowResponseSetter;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import io.undertow.server.HttpServerExchange;
import io.undertow.servlet.handlers.ServletRequestContext;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;

import javax.servlet.http.HttpServletResponse;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSHandlerInstrumentation implements TypeInstrumentation {

    @Override
    public ElementMatcher.Junction<ClassLoader> classLoaderOptimization() {
        return hasClassesNamed("io.undertow.servlet.handlers.FilterHandler");
    }

    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return named("io.undertow.servlet.handlers.FilterHandler");
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
                named("handleRequest")
                        .and(takesArgument(0, named("io.undertow.server.HttpServerExchange")))
                        .and(isPublic()),
                this.getClass().getName() + "$HandleRequestAdvice");
    }

    @SuppressWarnings("unused")
    public static class HandleRequestAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHandlerInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSUndertowHandlerCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            durationPrinter.printElapseDuration("HandleRequestAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Argument(value = 0, readOnly = false) HttpServerExchange exchange,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
            DurationPrinter durationPrinter = new DurationPrinter("XSHandlerInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if(callDepth.decrementAndGet() > 0) {
                return;
            }
            try {
                final ServletRequestContext servletRequestContext = exchange.getAttachment(ServletRequestContext.ATTACHMENT_KEY);
                XSBaggage baggageForCurrent = XSBaggage.fromContext(Java8Bridge.currentForXSContext());
                if(servletRequestContext != null && servletRequestContext.getServletResponse() != null && !baggageForCurrent.isEmpty()) {
                    HttpServletResponse httpServletResponse = (HttpServletResponse)servletRequestContext.getServletResponse();
                    W3CReentryBaggagePropagator.getInstance().inject(Java8Bridge.currentForXSContext(),httpServletResponse, UndertowResponseSetter.INSTANCE);
                    ReentryContextStorage.INSTANCE.close();
                }
            } catch (Throwable e) {
                StartConfig.throwableStackPrinter(e);
            }
            durationPrinter.printElapseDuration("HandleRequestAdvice e");
        }
    }

}
